import akshare as ak
import pandas as pd
from datetime import datetime
import os


# 按格式YYYYMMDD获取当前日期
def get_date():
    return datetime.now().strftime('%Y%m%d')


etf_symbols = {
    '510050': '华夏上证50ETF',
    '159915': '易方达创业板ETF',
    '512000': '券商ETF（SH）'
}

stock_symbols = {
    '603290': '斯达半导',
    '002049': '紫光国微',
    '603893': '瑞芯微',
    '603019': '中科曙光',
    '000066': '中国长城',
    '688981': '中芯国际',
    '000977': '浪潮信息',
    '300503': '昊志机电',
    '600745': '闻泰科技',
    '300059': '东方财富',
    '002156': '通富微电'
}

out_dir = './datas'
# 中文列名到英文名的映射（覆盖AKShare常见字段）
column_mapping = {
    "日期": "date",
    "开盘": "open",
    "开盘价": "open",  # 兼容不同函数返回的字段名
    "最高": "high",
    "最高价": "high",
    "最低": "low",
    "最低价": "low",
    "收盘": "close",
    "收盘价": "close",
    "成交量": "volume",
    "成交额": "amount",
    "涨跌幅": "pct_change",
    "涨跌额": "price_change",
    "换手率": "turnover_rate",
    "振幅": "amplitude",
    "股票代码": "code",
    "代码": "code",
    "名称": "name"
}


def get_real_quote():
    # 获取A股实时行情（深沪京A股）
    df = ak.stock_zh_a_spot()

    # 筛选出上证A股（代码以60开头）
    sh_df = df[df['代码'].str.startswith('60')]

    # 保存前100条数据
    sh_df.head(100).to_csv('A股实时行情_上证.csv', index=False)
    print(sh_df[['代码', '名称', '最新价', '涨跌幅']].head())


def get_etf_his_quote(symbol,
                      period='daily',
                      start_date="20200101",
                      end_date=get_date(),
                      adj='qfq'):
    # 获取宁德时代（300750）的历史行情
    df = ak.fund_etf_hist_em(
        symbol=symbol[0],
        # period=period,  # 周期：daily=日线，weekly=周线，monthly=月线
        start_date=start_date,
        end_date=end_date,
        adjust=adj  # 复权：qfq=前复权，hfq=后复权，""=不复权
    )

    if df is None or df.empty:
        print(f'从akshare获取{symbol[0]}_{symbol[1]}数据失败')
        return

    handle_result(df, symbol, period, start_date, end_date)


def get_stock_his_quote(symbol,
                        period='daily',
                        start_date="20200101",
                        end_date=get_date(),
                        adj='hfq'):
    # 获取宁德时代（300750）的历史行情
    df = ak.stock_zh_a_hist(
        symbol=symbol[0],
        period=period,  # 周期：daily=日线，weekly=周线，monthly=月线
        start_date=start_date,
        end_date=end_date,
        adjust=adj  # 复权：qfq=前复权，hfq=后复权，""=不复权
    )

    if df is None or df.empty:
        print(f'从akshare获取{symbol[0]}_{symbol[1]}数据失败')
        return

    handle_result(df, symbol, period, start_date, end_date)


def handle_result(df, symbol, period, start_date, end_date):
    # 重命名列名（只转换映射中存在的字段）
    df = df.rename(columns=column_mapping)
    # 如果‘code’列存在，则删除
    if 'code' in df.columns:
        df = df.drop(columns=['code'])
    df.to_csv(os.path.join(out_dir, f'{symbol[0]}_{symbol[1]}_{period}.csv'),
              index=False)
    print(
        f'{symbol[0]}_{symbol[1]}_{period}.csv 保存成功! start_date: {start_date}, end_date: {end_date}'
    )


if __name__ == "__main__":
    # get_real_quote()
    for symbol in etf_symbols.items():
        get_etf_his_quote(symbol)

    for symbol in stock_symbols.items():
        get_stock_his_quote(symbol)
